import base64
from io import BytesIO
from flask.helpers import make_response
import numpy as np

from app.api.transfer import image_utils
from . import api
from flask import json, request, jsonify
from uuid import uuid4
from .basic import ImageUtil
import re
from .transfer.stylizer import Stylizer
from .transfer import image_utils
s1 = Stylizer('./app/api/transfer/models/style1.pt')
s2 = Stylizer('./app/api/transfer/models/style2.pt')
s3 = Stylizer('./app/api/transfer/models/style3.pt')


@api.route('/health', methods=['GET'])
def health():
    return jsonify(code=0, msg='ok')

@api.route('/images/options', methods=['GET'])
def options():
    data = [
        {
            'kind': 1,
            'label': '图像灰度化'
        },
        {
            'kind': 2,
            'label': '图像二值化',
            'args': ['threshold', 'maxval']
        },
        {
            'kind': 3,
            'label': '旋转',
            'args': ['theta']
        },
        {
            'kind': 4,
            'label': '放缩',
            'args': ['fx', 'fy']
        },
        {
            'kind': 5,
            'label': '平移',
            'args': ['fx', 'fy']
        },
        {
            'kind': 6,
            'label': '傅里叶变换'
        },
        {
            'kind': 7,
            'label': '直方图'
        },
        {
            'kind': 8,
            'label': '直方图均衡化'
        },
        {
            'kind': 9,
            'label': '自适应阈值'
        },
        {
            'kind': 10,
            'label': '图像平滑',
            'args': ['kind'],
            'choice': ['gaussian', 'avg', 'median', 'bilateral']
        },
        {
            'kind': 11,
            'label': '形态学变换',
            'args': ['kind'],
            'choice': ['erosion', 'dilation', 'opening', 'closing', 'gradient', 'tophat', 'blackhat']
        },
        {
            'kind': 12,
            'label': '图像梯度',
            'args': ['kind'],
            'choice': ['laplacian', 'sobelx', 'sobely']
        },
        {
            'kind': 13,
            'label': '图像边缘'
        },
        {
            'kind': 14,
            'label': '霍夫变换',
            'args': ['kind'],
            'choice': ['normal', 'p'] #'circle'
        },
        {
            'kind': 15,
            'label': '语义分割'
        },
        {
            'kind': 16,
            'label': '图像去噪'
        }
    ]
    return jsonify(code=0, options=data)

@api.route('/images/processed', methods=['POST'])
def image_process():
    data = request.json
    kind = data.get('kind')
    image = data.get('image')
    args = data.get('args')
    result = process(image, kind, args)
    return jsonify(code=0, image=result)

@api.route('/images/transition', methods=['POST'])
def image_transition():
    data = request.json
    kind = data.get('kind')
    content_image = data.get('image')
    source = ImageUtil.to_bytes(content_image)
    # with open('./temp.jpg', 'wb') as f:
    #     f.write(source)
    image = image_utils.load(BytesIO(source))
    if kind == 1:
        stylized = s1.stylize(image)
    elif kind == 2:
        stylized = s2.stylize(image)
    elif kind == 3:
        stylized = s3.stylize(image)
    res = BytesIO()
    image_utils.save(stylized, res)
    return jsonify(code=0, image=ImageUtil.to_base64(res.getvalue()))

    
    

def process(image, kind, args):
    source = ImageUtil.from_array(image)
    if kind == 1:
        processed = ImageUtil.grayfy(source)
    elif kind == 2:
        if args:
            processed = ImageUtil.binarify(source, **args)
        else:
            processed = ImageUtil.binarify(source)
    elif kind == 3:
        if args:
            processed = ImageUtil.rotate(source, **args)
        else:
            processed = ImageUtil.rotate(source)
    elif kind == 4:
        if args:
            processed = ImageUtil.scale(source, **args)
        else:
            processed = ImageUtil.scale(source)
    elif kind == 5:
        if args:
            processed = ImageUtil.move(source, **args)
        else:
            processed = ImageUtil.scale(source)
    elif kind == 6:
        processed = ImageUtil.dft(source)
    elif kind == 7:
        processed = ImageUtil.hist(source)
    elif kind == 8:
        processed = ImageUtil.hist_equal(source)
    elif kind == 9:
        processed = ImageUtil.adaptive_threshold(source)
    elif kind == 10:
        if args:
            processed = ImageUtil.blurring(source, **args)
        else:
            processed = ImageUtil.blurring(source)
    elif kind == 11:
        if args:
            processed = ImageUtil.morphology(source, **args)
        else:
            processed = ImageUtil.morphology(source)
    elif kind == 12:
        if args:
            processed = ImageUtil.gradient(source, **args)
        else:
            processed = ImageUtil.gradient(source)
    elif kind == 13:
        processed = ImageUtil.canny_edge(source)
    elif kind == 14:
        if args:
            processed = ImageUtil.hough(source, **args)
        else:
            processed = ImageUtil.hough(source)
    elif kind == 15:
        processed = ImageUtil.segmentation(source)
    elif kind == 16:
        processed = ImageUtil.denoising(source)
    result = ImageUtil.to_base64(processed)
    return result